package mr

import (
	"log"
	"os"
	"sync"
	"time"
)
import "net"
import "net/rpc"
import "net/http"

type Coordinator struct {
	// Your definitions here.
	inputs             []string
	nReduce            int
	countDown          int
	mapWorkerStatus    []int //0: wait to do,1: wait,2:done
	mapMutex           sync.Mutex
	reduceWorkerStatus []int //0: wait to do,1: wait,2:done
	reduceMutex        sync.Mutex
	workStatus         int //0: need map, 1: need reduce, 2: done
	statusMutex        sync.Mutex

	finishCnt int
}

// Your code here -- RPC handlers for the worker to call.

//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

func (c *Coordinator) Register(args *RegisterWorkerArgs, reply *RegisterWorkerReply) error {
	//log.Printf("req register")
	status := 0
	c.statusMutex.Lock()
	status = c.workStatus
	c.statusMutex.Unlock()
	if status == 2 {
		reply.ReplyCode = 0
	} else {
		if status == 0 {
			c.assignMap(reply)
		} else if status == 1 {
			c.assignReduce(reply)
		}
	}
	return nil
}
func (c *Coordinator) assignMap(reply *RegisterWorkerReply) {
	taskID := -1
	done := true
	c.mapMutex.Lock()

	for i := 0; i < len(c.inputs); i++ {
		if c.mapWorkerStatus[i] != 2 {
			done = false
		}
		if c.mapWorkerStatus[i] == 0 {
			taskID = i
			c.mapWorkerStatus[i] = 1
			break
		}

	}

	c.mapMutex.Unlock()

	if done {
		c.assignReduce(reply)
	} else if taskID == -1 {
		reply.ReplyCode = 3
	} else {
		reply.ReplyCode = 1
		reply.TaskID = taskID
		reply.InputCnt = c.nReduce
		reply.FileNames = []string{c.inputs[taskID]}
		go func() {
			time.Sleep(time.Duration(c.countDown) * time.Second)
			c.mapMutex.Lock()
			if c.mapWorkerStatus[taskID] == 1 {
				c.mapWorkerStatus[taskID] = 0
			}
			c.mapMutex.Unlock()
		}()

		log.Printf("assigned map %d", reply.TaskID)
	}

}
func (c *Coordinator) assignReduce(reply *RegisterWorkerReply) {
	taskID := -1
	done := true
	c.reduceMutex.Lock()

	for i := 0; i < c.nReduce; i++ {
		if c.reduceWorkerStatus[i] != 2 {
			done = false
		}
		if c.reduceWorkerStatus[i] == 0 {
			taskID = i
			c.reduceWorkerStatus[i] = 1
			break
		}
	}

	c.reduceMutex.Unlock()
	if done {
		reply.ReplyCode = 0
	} else if taskID == -1 {
		reply.ReplyCode = 3
	} else {
		reply.ReplyCode = 2
		reply.TaskID = taskID
		reply.InputCnt = len(c.inputs)
		reply.FileNames = nil
		go func() {
			time.Sleep(time.Duration(c.countDown) * time.Second)
			c.reduceMutex.Lock()
			if c.reduceWorkerStatus[taskID] == 1 {
				c.reduceWorkerStatus[taskID] = 0
			}
			c.reduceMutex.Unlock()
		}()
		log.Printf("assigned reduce %d", reply.TaskID)
	}
}
func (c *Coordinator) WorkerReturn(args *WorkerRetArgs, reply *WorkerRetReply) error {
	if args.ErrCode != 0 {
		return nil
	}
	if args.WorkMode == 1 {
		c.mapMutex.Lock()
		c.mapWorkerStatus[args.TaskID] = 2
		c.mapMutex.Unlock()
	} else {
		c.reduceMutex.Lock()
		if c.reduceWorkerStatus[args.TaskID] != 2 {
			c.finishCnt++
			c.reduceWorkerStatus[args.TaskID] = 2
			if c.finishCnt == c.nReduce {
				c.statusMutex.Lock()
				c.workStatus = 2
				c.statusMutex.Unlock()
			}
		}
		c.reduceMutex.Unlock()
	}
	return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

//
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
//
func (c *Coordinator) Done() bool {
	ret := false
	status := 0
	c.statusMutex.Lock()
	status = c.workStatus
	c.statusMutex.Unlock()
	if status == 2 {
		ret = true
	}

	return ret
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}
	c.inputs = files
	c.nReduce = nReduce
	c.mapWorkerStatus = make([]int, len(files))
	c.workStatus = 0
	c.reduceWorkerStatus = make([]int, nReduce)
	c.countDown = 10
	c.finishCnt = 0
	c.server()
	log.Printf("Start Server")
	return &c
}
